//  Copyright (c) Microsoft Corporation.  All Rights Reserved.

using System;
using System.Collections.Generic;
using System.Text;
using System.Xml;
using System.IO;
using System.Security.Cryptography;
using System.Security.Cryptography.Xml;

using System.IdentityModel.Tokens;
using System.IdentityModel.Selectors;

using System.ServiceModel.Security.Tokens;

namespace EAI.SAML
{
    public class EAISamlEncryptedAssertion
    {
        EAISamlAssertion assertion;
        EncryptionParameters encryptionParameters;

        public EAISamlEncryptedAssertion(EAISamlAssertion assertion, EncryptionParameters encryptionParameters)
        {
            if (assertion == null)
                throw new ArgumentNullException("assertion");

            if (encryptionParameters == null)
                throw new ArgumentNullException("encryptionParameters");

            this.assertion = assertion;
            this.encryptionParameters = encryptionParameters;
        }

        public EAISamlEncryptedAssertion(XmlReader reader, EAISamlSerializer serializer, SecurityTokenSerializer keyInfoSerializer, SecurityTokenResolver outOfBandTokenResolver)
        {
            this.ReadXml(reader, serializer, keyInfoSerializer, outOfBandTokenResolver);
        }

        public EAISamlAssertion Assertion
        {
            get { return this.assertion; }
        }

        public virtual void ReadXml(XmlReader reader, EAISamlSerializer serializer, SecurityTokenSerializer keyInfoSerializer, SecurityTokenResolver outOfBandTokenResolver)
        {
            if (reader == null)
                throw new ArgumentNullException("reader");

            reader.MoveToContent();
            reader.Read();
            XmlDocument dom = new XmlDocument();
            dom.Load(reader);
            reader.ReadEndElement();

            SecurityKeyIdentifier decryptionTokenSki = null;
            SecurityToken decryptionToken = null;
            SymmetricSecurityKey decryptionKey = null;
            for (int i = 0; i < dom.DocumentElement.ChildNodes.Count; ++i)
            {
                if ((dom.DocumentElement.ChildNodes[i].LocalName == "KeyInfo") && (dom.DocumentElement.ChildNodes[i].NamespaceURI == "http://www.w3.org/2000/09/xmldsig#"))
                {
                    XmlDictionaryReader keyInfoReader = XmlDictionaryReader.CreateTextReader(Encoding.UTF8.GetBytes(dom.DocumentElement.ChildNodes[i].OuterXml), XmlDictionaryReaderQuotas.Max);
                    decryptionTokenSki = EAISamlSerializer.ReadSecurityKeyIdentifier(keyInfoReader, keyInfoSerializer);
                    decryptionToken = EAISamlSerializer.ResolveSecurityToken(decryptionTokenSki, outOfBandTokenResolver);
                    if (decryptionToken != null)
                        decryptionKey = (SymmetricSecurityKey)decryptionToken.SecurityKeys[0];
                    else
                        decryptionKey = (SymmetricSecurityKey)EAISamlSerializer.ResolveSecurityKey(decryptionTokenSki, outOfBandTokenResolver);

                    dom.DocumentElement.RemoveChild(dom.DocumentElement.ChildNodes[i]);
                    keyInfoReader.Close();
                    break;
                }
            }

            if (decryptionKey == null)
                throw new SecurityTokenException("Unable to get decryption key.");

            EncryptedData encData = new EncryptedData();
            encData.LoadXml(dom.DocumentElement);
            EncryptedXml encXml = new EncryptedXml();
            SymmetricAlgorithm symmetricAlgorithm = decryptionKey.GetSymmetricAlgorithm(encData.EncryptionMethod.KeyAlgorithm);
            byte[] decryptedBuffer = encXml.DecryptData(encData, symmetricAlgorithm);

            XmlDictionaryReader dicReader = XmlDictionaryReader.CreateTextReader(decryptedBuffer, XmlDictionaryReaderQuotas.Max);
            this.assertion = new EAISamlAssertion(dicReader, serializer, keyInfoSerializer, outOfBandTokenResolver);
            dicReader.Close();
        }

        public virtual void WriteXml(XmlWriter writer, SecurityTokenSerializer tokenSerializer)
        {
            MemoryStream ms = new MemoryStream();
            XmlDictionaryWriter dicWriter = XmlDictionaryWriter.CreateTextWriter(ms);
            this.assertion.WriteTo(dicWriter, tokenSerializer);
            dicWriter.Flush();
            ms.Position = 0;

            RNGCryptoServiceProvider rngCryptoServiceProvider = new RNGCryptoServiceProvider();
            byte[] crypto = new byte[this.encryptionParameters.EncryptionKeySize];
            rngCryptoServiceProvider.GetNonZeroBytes(crypto);
            WrappedKeySecurityToken wrappedToken = new WrappedKeySecurityToken(Guid.NewGuid().ToString(), crypto, this.encryptionParameters.EncryptedKeyAlgorithm, this.encryptionParameters.WrappingToken, this.encryptionParameters.WrappingTokenIdentifier);

            SecurityKeyIdentifier wrappingTokenIdentifier = new SecurityKeyIdentifier(new EncryptedKeyIdentifierClause(wrappedToken.GetWrappedKey(), wrappedToken.WrappingAlgorithm, wrappedToken.WrappingTokenReference));

            EncryptedXml encXml = new EncryptedXml();
            SymmetricAlgorithm symmetricAlgorithm = ((SymmetricSecurityKey)wrappedToken.SecurityKeys[0]).GetSymmetricAlgorithm(this.encryptionParameters.EncryptionAlgorithm);
            byte[] encryptedData = encXml.EncryptData(ms.ToArray(), symmetricAlgorithm);
            EncryptedData ed = new EncryptedData();
            ed.Type = EncryptedXml.XmlEncElementUrl;
            ed.EncryptionMethod = new EncryptionMethod(this.encryptionParameters.EncryptionAlgorithm);
            ed.CipherData.CipherValue = encryptedData;
            XmlElement encDataEle = ed.GetXml();

            dicWriter.Close();

            ms = new MemoryStream();
            dicWriter = XmlDictionaryWriter.CreateTextWriter(ms);
            tokenSerializer.WriteKeyIdentifier(dicWriter, wrappingTokenIdentifier);
            dicWriter.Flush();
            ms.Position = 0;
            XmlDocument ownerDocument = encDataEle.OwnerDocument;
            ownerDocument.Load(ms);
            encDataEle.AppendChild(ownerDocument.DocumentElement.CloneNode(true));
            dicWriter.Close();

            writer.WriteStartElement(EAISamlConstants.EncryptedAssertion, EAISamlConstants.Namespace);
            encDataEle.WriteTo(writer);
            writer.WriteEndElement();
        }

    }
}